# -*- coding: utf-8 -*-
"""
Created on Wed Jan 13 12:02:52 2021

@author: 26297
"""


#   模型预加载数据的预处理部分
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader
from age_net import AgeNet
import matplotlib.pyplot as plt
import torch.nn.functional as F
from time import time
import os
import numpy as np
from getImage2 import getImage
os.environ['CUDA_VISIBLE_DEVICE']='1'
import cv2
class AgeTrain:
    
    def __init__(self, model_path = '../model/MSE_model_001.pth'):
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')  # 处理器选择
        agenet = AgeNet.agenet(path)
        self.agenet = agenet.to(self.device)
        
        
        # 参数全部参与训练
        for param in self.agenet.parameters():
            param.requires_grad = True
        
        # 参数设置
        LR = 1e-6  # 学习率
        print('device:',self.device)
        # 平方差损失函数，计算二范数
        self.MSE_loss = nn.MSELoss()
        # 定义优化器
        self.optimizer = optim.Adam(self.agenet.parameters(), LR)
        
        
    
    def train(self, epoch, save_path, model_path, batch=1000,batch_size = 8, loss_txt = 'loss_record.txt', bias_txt = 'bias_record.txt'):
        if os.path.exists(model_path):
            print("存在已训练的模型，正在加载{}...".format(model_path))
            self.agenet.load_state_dict(torch.load(model_path))

        # 使用WeightRandomSampler来平衡数据集

        
        root = '../trainSet1'

        self.agenet.train()
        loss_list = []
        bias_sum = 0
        loss_sum = 0
        for i in range(1,batch+1):
            
            inputs,labels,_ = getImage(root_path = '../trainSet1', batch=batch_size)
            
            # 获得数据和对应的标签
            # inputs, labels = data
            inputs = inputs.to(self.device)
            labels = labels.to(self.device)
            inputs.requires_grad = True
            inputs = inputs.type(torch.cuda.FloatTensor)
            
            # 获得模型预测结果
            out = self.agenet(inputs).view(labels.size())
    
            # 计算二范数loss
            loss = self.MSE_loss(out.float(), labels.float())
            loss_list.append(loss.item())
            
            # 计算梯度
            loss.backward()
            # 修改权值
            self.optimizer.step()
            # 梯度清0
            self.optimizer.zero_grad()
            
            
            # 显示偏差和loss
            bias_list = []
            bias = abs(labels-out).sum().item()
            bias_sum += bias
            loss_sum += loss.item()
            times = 200
            if i % times == 0:
                print("第{}次训练".format(i), end='\t')
                # 显示loss和标签
                loss_aver = loss_sum / times
                labels_data = [i.item() for i in labels]
                pre_data = ['{0:.2f}'.format(i.item()) for i in out]
                print('loss_aver:{0:.4f} \n  label:{1}\n  pre:{2}'.format(loss_aver, labels_data, pre_data))
                
                # 显示偏差
                # bias = self.MSE_bias(train_loader)
                bias_aver = bias_sum/times/batch_size
                print('bias_aver:', bias_aver)
                bias_list.append(bias)
                bias_sum = 0  # 偏差值归零
                loss_sum = 0
                
                # with open(bias_txt, 'a+') as f:
                #     f.write(bias_aver)
                #     f.write('\n')

                with open(loss_txt,'a+') as f:
                    f.write('{0:.2f}'.format(loss_aver))
                    f.write('\n')
                    
                            
            if i>=batch:
                break
        
        
        # 展示loss变化 
        self.showChart(loss_list,epoch)
        self.showChart(bias_list, epoch)
        torch.save(self.agenet.state_dict(), save_path)
        return loss_list
        
    
    def test(self, model_path, batch_size = 8, bias_txt = 'bias_record.txt', class_acc_txt = 'class_record.txt'):
        try:
            self.agenet.load_state_dict(torch.load(model_path))
        except:
            print("模型不存在!")

        root = '../testSet'

        self.agenet.eval()
        loss_list = []
        bias_sum = 0
        loss_sum = 0
        count = 0
        image_count = np.zeros(10)
        image_acc = np.zeros(10)
        for i in range(100):
            
            # data = getImage(root_path = '../AFAD_image')
            
            # 获得数据和对应的标签
            inputs, labels, image_path = getImage(root_path='../testSet',batch=batch_size)
            inputs = inputs.to(self.device)
            labels = labels.to(self.device)
            inputs.requires_grad = False
            
            # 获得模型预测结果
            out = self.agenet(inputs).view(labels.size())     
            for i in range(len(image_path)):
                
                if abs(labels[i]-out[i]) <= 0.5:
                   image_acc[labels[i]] += 1
                   
                image_count[labels[i]] += 1 
                # plt.subplot(4, 6, i+1)
                # image =cv2.imread(image_path[i])
                # image = np.array(image)
                # plt.imshow(image[:,:, ::-1])
            #     print(i, labels[i].item(),out[i].item(),end='\n')
            # print(out)
            # plt.show()
            
            # 显示偏差和loss
            bias = abs(labels-out).sum().item()/batch_size
            bias_sum += bias*batch_size
            
            # if bias>1:
            #     count+=1
            #     print("labels:{}  pre:{}".format(labels, out))
        
        # 展示最后一次预测情况
        
        # 显示偏差
        bias_aver = bias_sum/batch_size/100
        print('bias_aver in test set:', bias_aver)
        # print("count",count)
        with open(bias_txt, 'a+') as f:
            f.write('{0:.2f}'.format(bias_aver))
            f.write('\n')

        acc_rate = [image_acc[i] / image_count[i] for i in range(len(image_count))]
        acc_rate = '\t'.join(['{0:.2f}'.format(acc_rate[i]) for i in range(len(acc_rate))])
        with open(class_acc_txt, 'a+') as f:
            f.write('{0}\n'.format(acc_rate))
        print(acc_rate)
    def showChart(self, data_list,ep):
        x = [i for i in range(len(data_list))]
        plt.title(ep)
        plt.plot(x, data_list)
        plt.show()

if __name__ == '__main__':
    print("模型加载...")
    path = '../Arcface_100.pth'
    
    loss_list = []  # 记录损失值的变化
    acc_list = []  # 记录准去率的变化
    start_epoch = 360
    end_epoch = 1000
    ageTrain = AgeTrain()
    # # 读取数据
    for ep in range(start_epoch, end_epoch):
        print("第{}个epoch:".format(ep))
        save_path = '../model/mse_{0:3d}'.format(ep)
        model_path = '../model/mse_{0:3d}'.format(ep-1)
        
        ageTrain.train(ep, save_path, model_path, batch=200,batch_size=24)
        ageTrain.test(save_path,batch_size=24)
    # model_path = '../model/tanh_{0:3d}'.format(start_epoch-1)
    # ageTrain.test(model_path)
        
    
        

        
